package kqutil

import (
	"context"
	"errors"
	"paper-pro/pkg/kqutil/kqbuffer"
	"sync"

	"github.com/segmentio/kafka-go"
	"github.com/zeromicro/go-zero/core/syncx"
)

var (
	ErrBusyRead      = errors.New("read buffer busy")
	ErrEmptyWrite    = errors.New("write buffer empty")
	ErrEmptyRead     = errors.New("read buffer empty")
	ErrCommitPending = errors.New("commit pending")
	ErrSignalStop    = errors.New("offset tracker op context.Context Done")
)

/*
 * 读写双缓冲区 kafka message 缓冲区
 * 支持单协程无锁写入；多协程并发获取 message 并标记已读；最后由单协程收集并提交offset
 */
type OffsetTracker struct {
	stopFlag   *syncx.AtomicBool
	ctx        context.Context
	cancelFunc func()
	bufferPool *kqbuffer.KqBufferPool

	writeBuffer kqbuffer.KqWriteBuffer

	rMutex       sync.Mutex // 读缓冲区保护锁
	readBuffer   kqbuffer.KqReadBuffer
	condWaitRead *sync.Cond

	maxRetryCommitTimes int // 最大连续 commit pending 时，重复尝试提交的次数
	alreadyRetryCommit  int
}

func NewOffsetTracker(capacity int, maxRetryCommitTimes int) *OffsetTracker {
	if capacity < 0 {
		return nil
	}
	ctx, cancel := context.WithCancel(context.Background())
	ot := &OffsetTracker{
		stopFlag:            syncx.ForAtomicBool(false),
		ctx:                 ctx,
		cancelFunc:          cancel,
		bufferPool:          kqbuffer.NewKqBufferPool(2, capacity),
		maxRetryCommitTimes: maxRetryCommitTimes,
		alreadyRetryCommit:  0,
	}
	ot.writeBuffer, _ = ot.bufferPool.GetEmptyWriteBuffer(context.Background())
	ot.readBuffer = nil
	ot.condWaitRead = sync.NewCond(&ot.rMutex)
	return ot
}

/*
 * OffsetTracker 是否处于环境终止状态
 */
func (ot *OffsetTracker) IsStop() bool {
	return ot.stopFlag.True()
}

/*
 * 通知所有阻塞操作终止
 */
func (ot *OffsetTracker) SignalStop() {
	if ot.stopFlag.CompareAndSwap(false, true) {
		ot.cancelFunc()
		ot.condWaitRead.Broadcast()
	}
}

/*
 * read buffer 是否已空
 * 内部用，使用前需加锁 rMutex
 */
func (ot *OffsetTracker) isEmptyRead() bool {
	return ot.readBuffer == nil || ot.readBuffer.IsEmpty()
}

/*
 * 设置当前 write buffer 为 read buffer，获取新的写缓冲区。
 * 非阻塞，若不成功直接返回设置结果
 */
func (ot *OffsetTracker) TrySetWriteToRead() (err error) {
	if newBuf, ok := ot.bufferPool.TryGetEmptyWriteBuffer(); ok {
		err = ot.setWriteToRead(newBuf)
		if err != nil {
			ot.bufferPool.PushBackReadBuffer(newBuf.AsRead())
		}
	} else {
		return ErrBusyRead // 读缓冲区还在使用，未被释放
	}

	return
}

func (ot *OffsetTracker) setWriteToRead(emptyBuffer kqbuffer.KqWriteBuffer) (err error) {
	if ot.writeBuffer.IsEmpty() {
		return ErrEmptyWrite
	}

	ot.rMutex.Lock()
	defer ot.rMutex.Unlock()

	if !ot.isEmptyRead() { // readBuffer 为空条件由 emptyBuffer 保证，即 emptyBuffer 为原来的 readBuffer
		return ErrBusyRead
	}

	rdBuffer := ot.writeBuffer.AsRead()
	ot.writeBuffer = emptyBuffer
	ot.readBuffer = rdBuffer

	ot.condWaitRead.Signal()

	return
}

/*
 * 消息写 write buffer。缓冲区满时阻塞
 */
func (ot *OffsetTracker) PushMsg(msg *kafka.Message) (isSetRead bool, err error) {
	if ot.writeBuffer.IsFull() {
		newBuf, err := ot.bufferPool.GetEmptyWriteBuffer(ot.ctx) // 等待读缓冲区为空，读写双缓冲区场景下，这里保证能成功
		if err != nil {
			if err == kqbuffer.ErrCtxDone {
				err = ErrSignalStop
			}
			return false, err
		}
		err = ot.setWriteToRead(newBuf)
		if err != nil {
			ot.bufferPool.PushBackReadBuffer(newBuf.AsRead())
		}
		isSetRead = true
	}

	return isSetRead, ot.writeBuffer.Push(msg)
}

/*
 * 获取批量 message。若获取为空，阻塞直到获取数据
 */
func (ot *OffsetTracker) BatchMsg(batchSize int) (batch *kqbuffer.KqBatchMsg, err error) {
	ot.rMutex.Lock()
	defer ot.rMutex.Unlock()

	for ot.isEmptyRead() || ot.readBuffer.IsBatchEmpty() {
		if ot.stopFlag.True() {
			return nil, ErrSignalStop
		}
		ot.condWaitRead.Wait()
	}

	return ot.readBuffer.Batch(batchSize)
}

/*
 * 标记批量 message 消费完毕，可以提交。返回剩余待处理完毕消息数量
 */
func (ot *OffsetTracker) ReadyMsg(batch *kqbuffer.KqBatchMsg) (countWaitReady int, err error) {
	ot.rMutex.Lock()
	defer ot.rMutex.Unlock()

	if ot.isEmptyRead() {
		return 0, ErrEmptyRead
	}

	return ot.readBuffer.Ready(batch), nil
}

/*
 * 收集各分区已消费消息的提交。使用完毕后释放 readBuffer
 */
func (ot *OffsetTracker) CollectOffset() (result kqbuffer.PartitionOffsetMap, err error) {
	ot.rMutex.Lock()
	defer ot.rMutex.Unlock()

	if ot.isEmptyRead() {
		return nil, kqbuffer.ErrPopEmpty
	}

	result, err = ot.readBuffer.Pop()
	if ot.readBuffer.IsEmpty() {
		var oldBuf kqbuffer.KqReadBuffer
		oldBuf, ot.readBuffer = ot.readBuffer, nil
		ot.bufferPool.PushBackReadBuffer(oldBuf)
	}

	switch err {
	case nil:
		ot.alreadyRetryCommit = 0
	case kqbuffer.ErrPopPending:
		ot.alreadyRetryCommit++
		if ot.alreadyRetryCommit >= ot.maxRetryCommitTimes {
			return nil, ErrCommitPending
		}
	}

	return
}
